package 中等.前缀和与差分;

import java.util.HashMap;
import java.util.Map;

/**
 * 给定一个整数数组 nums 和一个整数 k ，返回其中元素之和可被 k 整除的
 * （连续、非空） 子数组 的数目。
 * 子数组 是数组的 连续 部分。
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode-cn.com/problems/subarray-sums-divisible-by-k/
 */
public class 和可被K整除的子数组_974 {

    public static void main(String[] args) {

        int[] nums = {4, 5, 0, -2, -3, 1};
        System.out.println(subarraysDivByK(nums, 5));

    }

    /**
     * 前缀和求 子数组 [i,j] 的总和 == prefixSum[j]-prefixSum[i-1]
     * 求子数组是否可被 k 整除 等价于 (prefixSum[j]-prefixSum[i-1]) % k == 0
     * 根据同余定理
     * (prefixSum[j]-prefixSum[i-1]) % k == 0
     * -->  prefixSum[j] % k == prefixSum[i-1] % k
     * <p>
     * 将 prefix[i-1] % k 存进map，记录个数
     *
     * @param nums
     * @param k
     * @return
     */
    public static int subarraysDivByK(int[] nums, int k) {
        int count = 0;
        Map<Integer, Integer> record = new HashMap<>();
        // 注意: 当第一次 curSum % k == 0 时，[0,i] 是一个符合条件的子数组
        record.put(0, 1);
        int curSum = 0;
        for (int i = 0; i < nums.length; i++) {
            curSum += nums[i];
            // -7 % 3 = -1，-7 mod 3 = 2，需要把负数转为正数，curSum % k + k --> -1 + 3 = 2
            int mod = (curSum % k + k) % k;
            count += record.getOrDefault(mod, 0);
            record.put(mod, record.getOrDefault(mod, 0) + 1);
        }
        return count;
    }

}
